global numBins

classLabels = ['A' 'B' 'C' 'D' 'F' 'G' 'H' 'K' 'M' 'O' 'P' 'S' 'Z']; % note: currently do not use W - water bottle
C = numel(classLabels);
trainPath = './images/training/';
masksPath = './images/trainingmasks/';

testImgsPath = './images/test/easy/';
colorSpace = 'hsi';
similarity = 'weightedComponents';

numBins = 100;

% Build the training models
if ~exist(['./data/models_' colorSpace '_' num2str(numBins) '.mat'])
    models = -ones(5,numBins,C);
    for i=1:C
        classPath = [trainPath classLabels(i) '/'];
        list = dir([classPath '*.jpg']);
        list = {list(:).name};
        % For all training images of class i, build model color and texture
        % histograms
        Chist1 = zeros(1,numBins); Chist2 = zeros(1,numBins); Chist3 = zeros(1,numBins);
        Thist1 = zeros(1,numBins); Thist2 = zeros(1,numBins);
        for j=1:numel(list)
            imgName = list{j};
            I = imread([classPath imgName]);
            mask = load([masksPath imgName(1:end-4)]);
            
            [chist1,chist2,chist3] = Color_Features(I,mask.mask,colorSpace,false);
            [thist1,thist2] = Texture_Features(mask.mask,I,0);
            Chist1 = Chist1 + chist1; Chist2 = Chist2 + chist2; Chist3 = Chist3 + chist3;
            Thist1 = Thist1 + thist1; Thist2 = Thist2 + thist2;
        end
        % Normalize the class histograms
        models(1,:,i) = Chist1/sum(Chist1);
        models(2,:,i) = Chist2/sum(Chist2);
        models(3,:,i) = Chist3/sum(Chist3);
        models(4,:,i) = Thist1/sum(Thist1);
        models(5,:,i) = Thist2/sum(Thist2);
        disp(['Models built for class ' classLabels(i) '.'])
    end
    save(['./data/models_' colorSpace '_' num2str(numBins)], 'models')
else
    load(['./data/models_' colorSpace '_' num2str(numBins)])
end

% Build the training mean HS models to be used for segmentation
if ~exist('./data/models_meansHS.mat')
    models_meansHS = -ones(C,2);
    for i=1:C
        classPath = [trainPath classLabels(i) '/'];
        list = dir([classPath '*.jpg']);
        list = {list(:).name};
        
        sumPixels = 0;
        sumHue = 0;
        sumSat = 0;
        for j=1:numel(list)
            imgName = list{j};
            I = imread([classPath imgName]);
            mask = load([masksPath imgName(1:end-4)]); mask = mask.mask;
            hsv = rgb2hsv(I);
            H = hsv(:,:,1);
            S = hsv(:,:,2);
            sumPixels = sumPixels + sum(mask(:));
            sumHue = sumHue + sum(H(mask));
            sumSat = sumSat + sum(S(mask));
        end
        models_meansHS(i,1) = sumHue/sumPixels;
        models_meansHS(i,2) = sumSat/sumPixels;
        disp(['Mean hue and saturation saved for class ' classLabels(i) '.'])
    end
    save('./data/models_meansHS.mat', 'models_meansHS');
else
    load('./data/models_meansHS.mat')
end

% Make a mask of true classes for the test images
list = dir([testImgsPath '*.jpg']); list = {list(:).name};
T = numel(list); % T = number of test images
testClasses = zeros(1,T);
for i=1:numel(list)
    parts = textscan(list{i}, '%s%s%s', 'Delimiter', '-.');
    testClasses(i) = strfind(classLabels, char(parts{1}));
end
truthMask = zeros(C,T);
truthMask(sub2ind(size(truthMask), testClasses, 1:T)) = 1;


% Classify each test image
scores = zeros(C,T);
tic;
for t=1:T
    I = imread([testImgsPath list{t}]);
    
    % Get mask(s) from segmentation results
    ROIs = segmentationL(I, models_meansHS);
    
    ROIscores = zeros(C,size(ROIs,3));
    for r=1:size(ROIs,3)
        mask = ROIs(:,:,r);
        [chist1,chist2,chist3] = Color_Features(I,mask,colorSpace,true);
        [thist1,thist2] = Texture_Features(mask,I,1);
        
        % ----------------------------------------------------------------
        % Calculate similarity of test histogram(s) with the model
        % histogram(s) from each class and assign to class with highest
        % similarity
        
        % ----------------------------------------------------------------
    end
    
end
time = toc;

% Calculate accuracy
[~, classes] = sort(scores,'descend');

% Find misclassified images

